/*
  ngrams - построить частоту n-грамм
  суммарная частота всех n-грамм записывается в конце под ключом ' ' (один пробел)
*/

#include <cassert>

#include <iostream>
#include <sstream>
#include <fstream>
#include <map>

#include "../util.h"

using std::string;
using std::map;
using std::stringstream;

map<string, size_t> CalculateNGrams(size_t n, std::istream& in) {
    map<string, size_t> result;
    string line;
    while (getline(in, line)) {
        // добавить токены начала и конца предложения
        line = "<s> " + line + " </s>";
        stringstream ss(line);
        string token;
        string currentNGram = "";
        size_t currentNGramLength = 0;
        while (ss >> token) {
            if (currentNGram != "") {
                currentNGram += ' ';
            }
            currentNGram += token;
            ++currentNGramLength;
            if (currentNGramLength > n) {
                size_t firstSpacePos = currentNGram.find(' ');
                assert(firstSpacePos != string::npos);
                currentNGram = currentNGram.substr(firstSpacePos + 1);
                --currentNGramLength;
            }
            if (currentNGramLength == n) {
                ++result[currentNGram];
            }
        }
    }
    return result;
}

void PrintNGrams(const map<string, size_t>& ngrams, std::ostream& out) {
    size_t ngramsTotalFreq = 0;
    for (map<string, size_t>::const_iterator it = ngrams.begin(); it != ngrams.end(); ++it) {
        out << it->first << '\t' << it->second << '\n';
        ngramsTotalFreq += it->second;
    }
    // Последней строкой под ключом ' ' (один пробел) записывается общая сумма встречаемостей n-грамм.
    out << " \t" << ngramsTotalFreq << '\n';
}

void PrintUsage() {
    std::cerr << "Usage: ngrams N input_file [output_file], 1 <= count <= 5\n"
              << "If output_file isn't given the result is printed to stdout\n";
}

int main(int argc, char **argv) {
    if ((argc < 3) || (argc > 4)) {
        PrintUsage();
        return 1;
    }

    size_t n = FromString<size_t>(argv[1]);
    if (n < 1 || n > 5) {
        PrintUsage();
        return 1;
    }

    std::ifstream input(argv[2]);
    map<string, size_t> ngrams = CalculateNGrams(n, input);
    
    if (argc == 4) {
        std::ofstream output(argv[3]);
        PrintNGrams(ngrams, output);
    } else {
        assert(argc == 3);
        PrintNGrams(ngrams, std::cout);
    }

    return 0;
}
